from mindspore import ops
from mindspore_gs.ptq.ptq.wrapper_cell import WrapperCell

class AllQuantLowRankInferCell(WrapperCell):
    def __init__(self,
                 layer_name: str,
                 q0: AllQuantLinearInferCell,
                 q1: AllQuantLinearInferCell):
        # 调用父类构造，不需要原始 seq，后面直接用 inf0/inf1
        super().__init__(layer_name, None, None, None)

        # 第一次完整量化段
        self.inf0 = q0.deploy()

        # 第二次量化段，拆出内部算子
        inf1 = q1.deploy()
        self.inf1 = inf1

        # 按照 dir 输出选出的算子和 bias
        self.act_quant_op = inf1.layer.activation   # float→int8
        self.matmul_op    = inf1.layer.matmul       # int8×int8→int8
        self.bias_add     = inf1.layer.bias_add     # int8 bias 加法
        self.dequant_op   = inf1.layer.to_float     # int8→float

    def construct(self, x):
        # —— 第一次量化 MatMul ——  
        x_mid = self.inf0(x)  

        # —— 第二次量化 MatMul ——  
        x_q  = self.act_quant_op(x_mid)            # 1) 激活量化
        y_q  = self.matmul_op(x_q)                 # 2) MatMul
        y_q  = self.bias_add(y_q, self.inf1.layer.bias)  # 3) 加 bias
        y    = self.dequant_op(y_q)                # 4) 反量化
        return y
